#!/usr/bin/env python3
'''
未验证
通过训练可以达到900+分，接近1000分，待play验证
可能需要调整学习率

将学习率下调到1e-4,发现成绩是可以更加稳定在满分附近（不调整也可以有时候接近满分1000分），但是依旧无法达到满分，可能还需要下调
也有可能这是遗传算法的问题
'''
import gymnasium as gym
import ptan
import time
import argparse
import numpy as np
import collections
import os
import torch
import torch.nn as nn
from torch import multiprocessing as mp
from torch import optim

from tensorboardX import SummaryWriter

from lib import common

NOISE_STD = 0.05
LEARNING_RATE = 0.01
PROCESSES_COUNT = 6
ITERS_PER_UPDATE = 10
MAX_ITERS = 100000

# result item from the worker to master. Fields:
# 1. random seed used to generate noise
# 2. reward obtained from the positive noise
# 3. reward obtained from the negative noise
# 4. total amount of steps done
RewardsItem = collections.namedtuple('RewardsItem', field_names=['seed', 'pos_reward', 'neg_reward', 'steps'])


def make_env():
    return gym.make("InvertedPendulum-v4")


class Net(nn.Module):
    def __init__(self, obs_size, act_size, hid_size=64):
        super(Net, self).__init__()
        # 与前一章节不同的地方在于，这里输出的激活函数使用的是tanh
        # 因为游戏环境变成了连续值的动作
        # 输出的是中值
        self.mu = nn.Sequential(
            nn.Linear(obs_size, hid_size),
            nn.Tanh(),
            nn.Linear(hid_size, hid_size),
            nn.Tanh(),
            nn.Linear(hid_size, act_size),
            nn.Tanh(),
        )

    def forward(self, x):
        x = x.float()
        return self.mu(x) * 3


def evaluate(env, net, device="cpu"):
    obs, _ = env.reset()
    reward = 0.0
    steps = 0
    while True:
        obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
        action_v = net(obs_v)
        action = action_v.data.cpu().numpy()[0]
        obs, r, done, trunc, _ = env.step(action)
        done = done or trunc
        reward += r
        steps += 1
        if done:
            break
    return reward, steps


def sample_noise(net, device="cpu"):
    res = []
    neg = []
    for p in net.parameters():
        noise_t = torch.FloatTensor(np.random.normal(size=p.data.size()).astype(np.float32)).to(device)
        res.append(noise_t)
        neg.append(-noise_t)
    return res, neg


def eval_with_noise(env, net, noise, noise_std, device="cpu"):
    for p, p_n in zip(net.parameters(), noise):
        p.data += noise_std * p_n
    r, s = evaluate(env, net, device)
    for p, p_n in zip(net.parameters(), noise):
        p.data -= noise_std * p_n
    return r, s


def compute_ranks(x):
    """
    返回x展平后的每个值的在整体的顺序
    Returns ranks in [0, len(x))
    Note: This is different from scipy.stats.rankdata, which returns ranks in [1, len(x)].
    """
    # 检测张量的维度是否为1
    assert x.ndim == 1
    # 创建一个长度和x一样的空numpy矩阵
    ranks = np.empty(len(x), dtype=int)
    # 按照x小到大排序后的索引，对对应索引的值，按顺序填入0,1,2...
    ranks[x.argsort()] = np.arange(len(x))
    return ranks


def compute_centered_ranks(x):
    '''
    return 计算x中每个元素的值在整体的排名索引，再归一化到-0.5~0.5之间
    维度和x一致
    '''

    # ravel展平张量维度为一维，并切是引用模式，所以修改ravel后的值，会影响原来的值
    # compute_ranks(x.ravel())得到x中每个值的排名
    # 将排名还原为x的维度，类型为float32
    y = compute_ranks(x.ravel()).reshape(x.shape).astype(np.float32)
    # 将排名归一化到-0.5~0.5之间
    y /= (x.size - 1)
    y -= .5
    return y


def train_step(optimizer, net, batch_noise, batch_reward, writer, step_idx, noise_std):
    # 保存每个噪音的权重，权重的大小按照奖励的大小进行
    weighted_noise = None
    # 得到bacth_reward中每个奖励值从小到大的名次，且是归一化到-0.5~0.5之间
    # 用名字来表示奖励的大小，这样可以更好的进行权重的计算
    # 这里之所以采用名字，是因为在这里的游戏中，奖励是有负值的
    # 为了避免负值对权重更新影响大，所以这里采用了排名的方式，也间接实现了
    # 奖励大的影响权重大，奖励小的影响权重小，成功应用了进化策略算法
    norm_reward = compute_centered_ranks(np.array(batch_reward))

    for noise, reward in zip(batch_noise, norm_reward):
        if weighted_noise is None:
            weighted_noise = [reward * p_n for p_n in noise]
        else:
            # 不断的将噪音和奖励权重相乘法后，累加到weighted_noise中
            for w_n, p_n in zip(weighted_noise, noise):
                w_n += reward * p_n
    m_updates = [] #
    optimizer.zero_grad()
    for p, p_update in zip(net.parameters(), weighted_noise):
        # 为什么要除以noise_std方差
        # 噪声标准差（noise_std）：噪声标准差是我们加到参数上的随机噪声的标准差。在进化策略中，这个噪声用来探索参数空间。噪声的大小（标准差）会影响参数变化的幅度。为了确保参数更新不依赖于噪声的具体尺度，我们需要通过噪声标准差来规范化更新。这样，无论噪声标准差的绝对大小如何，更新步骤都保持一致。
        # 样本数量：在计算梯度时，我们通常对样本进行平均，以得到对整个数据集的一个泛化估计。这里，我们使用的是基于一组样本的估计，所以需要除以样本数量，以得到平均效应。
        # 模拟学习率，如果噪声越大，那么更新的幅度就越大
        #
        update = p_update / (len(batch_reward) * noise_std)
        # 为什么这里要用减？要注意，这里更新的是梯度，所以没有学习率
        # P330
        # 在强化学习和某些其他类型的优化问题中，我们的目标是最大化某个指标而不是最小化。通常，在神经网络训练和许多传统优化任务中，我们关注的是最小化损失函数，这就是所谓的梯度下降。但在强化学习中，特别是当我们的目标是最大化总奖励时，这就转变成了一个梯度上升问题。
        #
        # ### 梯度上升的目的
        #
        # 在强化学习中，我们通常想要最大化智能体在给定任务中获得的总奖励。为了实现这一点，我们需要找到能够产生最高奖励的策略参数。这就是为什么我们需要梯度上升：
        #
        # 1. **最大化奖励**：我们的目标是找到能够使奖励最大化的参数设置。在数学上，这意味着我们要找到目标函数（例如，预期总奖励）的最大值。
        #
        # 2. **梯度方向**：梯度指向函数增长最快的方向。在最小化问题（如损失最小化）中，我们沿着梯度的反方向（梯度下降）移动参数。但在最大化问题中，我们要沿着梯度的正方向（梯度上升）移动参数，以增加函数的值。
        #
        # ### 梯度上升的实现
        #
        # 在实现上，梯度上升可以看作是梯度下降的“镜像”。如果你有一个关于参数 \( \theta \) 的目标函数 \( J(\theta) \)（在强化学习中通常是累积奖励），那么通过对 \( J \) 求梯度并朝着梯度的方向调整 \( \theta \)，你可以增加 \( J \) 的值。
        #
        # 在代码中，如果你使用的是一个标准的优化器（这些通常设计为执行梯度下降），那么你可以通过对梯度取负来实现梯度上升。也就是说，即使优化器本身在尝试减小一个值，通过输入负梯度，实际上它会增加你的目标函数。
        #
        # ### 结论
        #
        # 总之，梯度上升在强化学习和某些优化问题中至关重要，因为我们的目标是最大化（而不是最小化）一个函数。这通常涉及到找到使得奖励或某个评价指标最大化的参数，而梯度上升正是实现这一目标的有效方法。
        # ===========================
        # 这行代码将计算出的更新量设置为参数的负梯度。这样做的原因是在优化过程中，我们通常执行梯度下降（尽管这里使用的是梯度上升来最大化奖励）。通过设置 p.grad 为 -update，当调用优化器的 step() 方法时，它实际上会对参数执行 p += -lr * (-update) = p + lr * update，即朝着提高奖励的方向调整参数。
        # 所以这里采用这种方式就是因为优化器的更新网络参数的方式导致的
        # 所以这里本质上适和01方式是一样的
        p.grad = -update
        m_updates.append(torch.norm(update))
    # 记录更新量的l2范数的平局值
    writer.add_scalar("update_l2", np.mean(m_updates), step_idx)
    optimizer.step()


def worker_func(worker_id, params_queue, rewards_queue, device, noise_std):
    '''
    子进程
    '''
    # 创建环境
    env = make_env()
    # 创建网络
    net = Net(env.observation_space.shape[0], env.action_space.shape[0]).to(device)
    net.eval()

    while True:
        params = params_queue.get()
        if params is None:
            break
        # 加载得到的网络参数
        net.load_state_dict(params)

        # 每ITERS_PER_UPDATE轮次更新一次网络参数
        for _ in range(ITERS_PER_UPDATE):
            seed = np.random.randint(low=0, high=65535)
            np.random.seed(seed)
            # 同样也是采样
            noise, neg_noise = sample_noise(net, device=device)
            # 同样也是评估拿到奖励
            pos_reward, pos_steps = eval_with_noise(env, net, noise, noise_std, device=device)
            neg_reward, neg_steps = eval_with_noise(env, net, neg_noise, noise_std, device=device)
            # 传递种子就是为了提高进程的通信效率
            rewards_queue.put(RewardsItem(seed=seed, pos_reward=pos_reward,
                                          neg_reward=neg_reward, steps=pos_steps+neg_steps))
    pass


if __name__ == "__main__":
    # 理解这里采用了什么多进程模式
    # fork启动方式：
    #     子进程会直接使用主进程的资源，但是不会重新执行主进程的代码（if __name__='__main__':之外的代码），子进程会直接执行if __name__='__main__':内的代码，此处的代码为主进程创建子进程要执行的函数
    #
    # spawn启动方式：
    #     子进程不会直接使用主进程的资源，而是重新执行主进程的代码（if __name__='__main__':之外的代码），然后继续执行if __name__='__main__':内的代码，此处的代码为主进程创建子进程要执行的函数
    #
    # forkserver启动方式：
    #     先通过spawn启动forkserver进程，然后通过fork启动子进程
    #     forkserver是一个独立的进程，此后需要产生子进程的时候，父进程需要联系该进程 fork 一个子进程。因为 forkserver 本身是一个单线程进程，所以是线程安全的。而且，与 spawn 类似，子进程只会继承必要的资源。
    # 因为需要获取main函数以外的代码资源，所以这里采用了spawn方式启动子进程
    mp.set_start_method('spawn')
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action='store_true', help="Enable CUDA mode")
    parser.add_argument("--lr", type=float, default=LEARNING_RATE)
    parser.add_argument("--noise-std", type=float, default=NOISE_STD)
    parser.add_argument("--iters", type=int, default=MAX_ITERS)
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    args = parser.parse_args()
    device = "cuda" if args.cuda else "cpu"

    # 创建保存训练过程中数据的路径
    save_path = os.path.join("saves", "es-" + args.name)
    os.makedirs(save_path, exist_ok=True)

    writer = SummaryWriter(comment="-cheetah-es_lr=%.3e_sigma=%.3e" % (args.lr, args.noise_std))
    env = make_env()
    net = Net(env.observation_space.shape[0], env.action_space.shape[0])
    print(net)

    # 进程间的参数队列，用来同步神经网络的参数到子进程中
    params_queues = [mp.Queue(maxsize=1) for _ in range(PROCESSES_COUNT)]
    # 进程间的奖励队列，用来传递每个进程生成的奖励的
    rewards_queue = mp.Queue(maxsize=ITERS_PER_UPDATE)
    workers = []

    for idx, params_queue in enumerate(params_queues):
        # 创建并启动进程,传递的参数分别有：进程id，参数队列，奖励队列，设备，噪声方差（用于探索）
        proc = mp.Process(target=worker_func, args=(idx, params_queue, rewards_queue, device, args.noise_std))
        proc.start()
        workers.append(proc)

    print("All started!")
    # 不同这里创建优化器
    # 为了展示另外一种方式通过遗传算法更新网络权重
    optimizer = optim.Adam(net.parameters(), lr=args.lr)

    start_idx = 0
    # 增加加载模型的代码
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_idx = checkpoint['start_idx']
        print("加载模型成功")

    for step_idx in range(start_idx, args.iters):
        # broadcasting network params
        # 将神经网络的参数广播到每一个进程中
        start_idx = step_idx
        params = net.state_dict()
        for q in params_queues:
            q.put(params)

        # waiting for results
        t_start = time.time()
        batch_noise = []
        batch_reward = []
        results = 0
        batch_steps = 0
        batch_steps_data = []
        while True:
            while not rewards_queue.empty():
                # 如果进程有反馈奖励，则获取进程反馈的奖励
                # 进程所使用的随机种子也会负载奖励中
                reward = rewards_queue.get_nowait()
                # 设置该奖励所使用的随机种子，然后进行噪音采样
                # 由于是伪随机数，所以设置了相同的种子后，即可得到
                # 子进程中相同的采样正负噪音
                np.random.seed(reward.seed)
                noise, neg_noise = sample_noise(net)
                batch_noise.append(noise)
                batch_reward.append(reward.pos_reward)
                batch_noise.append(neg_noise)
                batch_reward.append(reward.neg_reward)
                results += 1
                batch_steps += reward.steps
                batch_steps_data.append(reward.steps)

            # 从子进程移植获取奖励知道奖励的采样数达到了进程数*每个进程采样数
            if results == PROCESSES_COUNT * ITERS_PER_UPDATE:
                break
            time.sleep(0.01)

        # 这个应该只是用于记录时间
        dt_data = time.time() - t_start
        # 计算所有奖励的平均值
        m_reward = np.mean(batch_reward)
        # 开始调整网络参数
        train_step(optimizer, net, batch_noise, batch_reward, writer, step_idx, args.noise_std)
        # 保存网络参数
        checkpoint = {
            "net": net.state_dict(),
            "optimizer": optimizer.state_dict(),
            "start_idx": start_idx
        }
        common.save_checkpoints(start_idx, checkpoint, save_path, "es", keep_last=5)
        # 分别记录奖励的平均值、方差、最大值、长度、 batch_steps（本次训练游戏能够进行的步数）、采样速度（本次训练的总步数以及所花费的时间）
        writer.add_scalar("reward_mean", m_reward, step_idx)
        writer.add_scalar("reward_std", np.std(batch_reward), step_idx)
        writer.add_scalar("reward_max", np.max(batch_reward), step_idx)
        writer.add_scalar("batch_episodes", len(batch_reward), step_idx)
        writer.add_scalar("batch_steps", batch_steps, step_idx)
        speed = batch_steps / (time.time() - t_start)
        writer.add_scalar("speed", speed, step_idx)
        dt_step = time.time() - t_start - dt_data

        print("%d: reward=%.2f, speed=%.2f f/s, data_gather=%.3f, train=%.3f, steps_mean=%.2f, min=%.2f, max=%.2f, steps_std=%.2f" % (
            step_idx, m_reward, speed, dt_data, dt_step, np.mean(batch_steps_data),
            np.min(batch_steps_data), np.max(batch_steps_data), np.std(batch_steps_data)))

    # 训练完成后，清空参数和等待进程结束
    for worker, p_queue in zip(workers, params_queues):
        p_queue.put(None)
        worker.join()
